# -*-coding:utf-8 -*-

"""
@File       : elastic_search.py
@Author     : HW Shen
@Date       : 2021/1/5
@Desc       :
"""

import arrow
import pytz
from elasticsearch import Elasticsearch
from elasticsearch.helpers import scan, streaming_bulk

es_zoo_connection = Elasticsearch(
    'http://esreader:NzdmMGM4MzE2@10.0.4.100:4000',
    dead_timeout=120,
    timeout=300,
    retry_on_timeout=True
)
es_yq_connection = Elasticsearch(
    'http://esreader:NzdmMGM4MzE2@10.0.4.100:5000',
    dead_timeout=10,
    retry_on_timeout=True
)


class ElasticSearch(object):

    client = es_zoo_connection
    index = None
    doc_type = None
    id_field = '_id'
    version = ''
    source_id_field = ''
    alias_field = ''
    alias_separator = ''
    alias_func = None
    tz = pytz.timezone('Asia/Shanghai')

    @classmethod
    def upsert(cls, doc_id, doc, index=None, doc_as_upsert=True, **kwargs):
        index = index or cls.index
        body = {
            'doc': doc,
        }
        if doc_as_upsert:
            body['doc_as_upsert'] = True

        cls.client.update(index, doc_id, body, doc_type=cls.doc_type, **kwargs)

    @classmethod
    def search(cls, query, index=None, source=False, **kwargs):
        index = index or cls.index
        search_data = cls.client.search(body=query, index=index, **kwargs)
        if source:
            search_data = [i['_source'] for i in search_data['hits']['hits']]
        return search_data

    @classmethod
    def scan(cls, query, index=None, **kwargs):
        index = index or cls.index
        return scan(cls.client, query=query, index=index, **kwargs)

    @classmethod
    def delete(cls, doc_id, index=None, **kwargs):
        index = index or cls.index
        cls.client.delete(index, doc_id, **kwargs)

    @classmethod
    def put_mapping(cls, body, index=None, **kwargs):
        """
        >>> body = {"properties": {"ocr": {"type": "text"}}}
        >>> ES.put_mapping(body, index='post_douyin')
        """
        index = index or cls.index
        cls.client.indices.put_mapping(body, index=index, **kwargs)

    @classmethod
    def put_template(cls, name, body):
        cls.client.indices.put_template(name, body)

    @classmethod
    def scan_all(cls, *args, index=None, **kwargs):
        s = cls.scan(*args, index=index, **kwargs)
        items = []
        cnt = 0
        for i in s:
            items.append(i['_source'])
            cnt += 1
            # if cnt % 1000 == 0:
            #     print('scan count: {}'.format(cnt))
        return items

    @classmethod
    def index_name(cls, doc):
        if cls.alias_field and cls.alias_field in doc.keys():
            alias_part = doc[cls.alias_field]
            if cls.version:
                index = '{}{}{}{}{}'.format(cls.index, cls.alias_separator, cls.version, cls.alias_separator,
                                            cls.alias_func(alias_part))
            else:
                index = '{}{}{}'.format(cls.index, cls.alias_separator, cls.alias_func(alias_part))
        else:
            index = cls.index
        return index

    @classmethod
    def id_name(cls, doc):
        _id = doc.get(cls.id_field) and doc.pop(cls.id_field) or doc.get(cls.source_id_field)
        # if not _id:
        #     print('========', doc)
        assert _id, 'doc _id must not be None'
        return _id

    @classmethod
    def bulk_upsert(cls, docs, **kwargs):
        """批量操作文章, 仅支持 index 和 update"""
        op_type = kwargs.get('op_type', 'update')
        upsert = kwargs.get('upsert', True)
        actions = cls._gen_bulk_actions(
            docs,
            cls.index_name,
            cls.doc_type,
            cls.id_name,
            op_type,
            upsert=upsert
        )
        result = streaming_bulk(
            cls.client,
            actions,
            chunk_size=3000,
            raise_on_error=False,
            raise_on_exception=False,
            max_retries=5,
            request_timeout=60
        )
        return result

    @classmethod
    def _gen_bulk_actions(
            cls,
            docs,
            index_name,
            doc_type,
            id_name,
            op_type,
            upsert=True,
            **kwargs
    ):
        assert not upsert or (upsert and op_type == 'update'), 'upsert should use "update" as op_type'

        for doc in docs:
            # 支持 index_name 作为一个工厂函数
            if callable(index_name):
                index = index_name(doc)
            else:
                index = index_name

            if op_type == 'index':
                _source = doc
            elif op_type == 'update' and not upsert:
                _source = {'doc': doc}
            elif op_type == 'update' and upsert:
                _source = {'doc': doc, 'doc_as_upsert': True}
            else:
                continue

            if callable(id_name):
                id = id_name(doc)
            else:
                id = id_name

            # 生成 Bulk 动作
            action = {
                "_op_type": op_type,
                "_index": index,
                "_type": doc_type,
                "_id": id,
                "_source": _source
            }
            yield action


class JoinPost(ElasticSearch):
    index = 'post'
    doc_type = '_doc'
    id_field = '_id'
    version = 'v1_1'
    alias_separator = '-'
    alias_field = 'post_create_time'
    alias_func = lambda x: '%d.%02d' % (x.year, x.month)
    join_field = 'join_field'
    join_relation = {
        'parent': 'post',
        'child': 'comment'
    }
    join_type = None

    @classmethod
    def id_name(cls, doc):
        doc_id = super().id_name(doc)
        return '{}_{}'.format(doc['platform'], doc_id)

    @classmethod
    def index_name(cls, doc):
        if cls.alias_field and doc[cls.alias_field]:
            aliase_part = doc[cls.alias_field]
            if isinstance(aliase_part, str):
                aliase_part = arrow.get(aliase_part)
            if isinstance(aliase_part, int):
                aliase_part = arrow.get(aliase_part).astimezone(cls.tz)
            if cls.version:
                index = '{}{}{}{}{}'.format(cls.index, cls.alias_separator, cls.version, cls.alias_separator,
                                            cls.alias_func(aliase_part))
            else:
                index = '{}{}{}'.format(cls.index, cls.alias_separator, cls.alias_func(aliase_part))
        else:
            index = cls.index
        return index

    @classmethod
    def _gen_bulk_actions(cls, docs, index_name, doc_type, id_name, op_type, upsert=True, **kwargs):
        assert not upsert or (upsert and op_type == 'update'), 'upsert should use "update" as op_type'

        for doc in docs:
            # 支持 index_name 作为一个工厂函数
            if callable(index_name):
                index = index_name(doc)
            else:
                index = index_name

            if callable(id_name):
                id = id_name(doc)
            else:
                id = id_name

            if op_type == 'index':
                _source = doc
            elif op_type == 'update' and not upsert:
                _source = {'doc': doc}
            elif op_type == 'update' and upsert:
                _source = {'doc': doc, 'doc_as_upsert': True}
            else:
                continue

            # 生成 Bulk 动作
            action = {
                "_op_type": op_type,
                "_index": index,
                "_type": doc_type,
                "_id": id,
                "_source": _source
            }
            yield action

    @classmethod
    def scan(cls, query, index=None, **kwargs):
        return scan(cls.client, query=query, index=index or cls.index, **kwargs)


class EsPost(JoinPost):
    join_type = 'parent'
    source_id_field = 'post_id'
    client = es_yq_connection

    @classmethod
    def _gen_bulk_actions(cls, docs, index_name, doc_type, id_name, op_type, upsert=True, **kwargs):
        assert not upsert or (upsert and op_type == 'update'), 'upsert should use "update" as op_type'

        for doc in docs:
            # 支持 index_name 作为一个工厂函数
            if callable(index_name):
                index = index_name(doc)
            else:
                index = index_name

            if callable(id_name):
                id = id_name(doc)
            else:
                id = id_name

            doc[cls.join_field] = {
                'name': cls.join_relation['parent']
            }

            if op_type == 'index':
                _source = doc
            elif op_type == 'update' and not upsert:
                _source = {'doc': doc}
            elif op_type == 'update' and upsert:
                _source = {'doc': doc, 'doc_as_upsert': True}
            else:
                continue

            # 生成 Bulk 动作
            action = {
                "_op_type": op_type,
                "_index": index,
                "_type": doc_type,
                "_id": id,
                "_source": _source
            }
            yield action


def put_mapping_by_alias():
    post_indices = es_zoo_connection.indices.get_alias("post_douyin*")
    for i in post_indices:
        response = es_zoo_connection.indices.put_mapping(
            {
                "properties": {
                    "ocr": {
                        "dynamic": "strict",
                        "properties": {
                            "text": {
                                "type": "text"
                            },
                            "frame": {
                                "type": "nested",
                                "dynamic": "strict",
                                "properties": {
                                    "time": {
                                        "type": "integer"
                                    },
                                    "text": {
                                        "type": "text"
                                    }
                                }
                            }
                        }
                    }
                }
            },
            index=i
        )
        # print(response)


class ESZoo(ElasticSearch):
    client = es_zoo_connection


class ESYq(ElasticSearch):
    client = es_yq_connection


class ESUserDou(ESZoo):
    index = 'user_douyin'


class ESLogsUserDou(ESZoo):
    index = 'logs_user_douyin'


class ESPostDou(ESZoo):
    index = 'post_douyin'


class ESStarDou(ESZoo):
    index = 'star_douyin'


class ESShopDou(ESZoo):
    index = 'shop_douyin'


class ESLogStarDou(ESZoo):       #注意：logs_star_douyin   表有问题，必须加上时间logs_star_douyin_2020
    index = 'logs_star_douyin_20*'


class ESLogStarDou202(ESZoo):
    index = 'logs_star_douyin_202*'

class ESLogStarDou201(ESZoo):
    index = 'logs_star_douyin_201*'

class ESLogUserDou(ESZoo):
    index = 'logs_user_douyin'


class ESLogShopDou(ESZoo):
    index = 'logs_shop_douyin_2020'


class ESFansDou(ESZoo):
    index = 'fans_douyin'


class ESLikeDou(ESZoo):
    index = 'like_douyin'


class ESChaDou(ESZoo):
    index = 'challenge_douyin'


class ESRoomDou(ESZoo):
    index = 'room_douyin'


class ESCommentDou(ESZoo):
    index = 'comment_douyin'


class ESLiveDou(ESZoo):
    index = 'live_douyin'


# 快手
class ESUserKwai(ESZoo):
    index = 'user_kwai'


class ESPostKwai(ESZoo):
    index = 'post_kwai'


class ESCommentKwai(ESZoo):
    index = 'comment_kwai'


class ESStarKwai(ESZoo):
    index = 'star_kwai'


class ESShopKwai(ESZoo):
    index = 'shop_kwai'


class ESLogShopKwai2020(ESZoo):
    index = 'logs_shop_kwai_2020'

# 小红书
class ESGoodsRed(ESZoo):
    index = 'goods_xiaohongshu'


class ESPostRed(ESZoo):
    index = 'post_xiaohongshu'


class ESUserRed(ESZoo):
    index = 'user_xiaohongshu'


class ESStarRed(ESZoo):
    index = 'star_xiaohongshu'


class ESLogsUserRed(ESZoo):
    index = 'logs_kol_xiaohongshu'


class ESLogsStarRed(ESZoo):
    index = 'logs_star_xiaohongshu_2020'



class ESCommentRed(ESZoo):
    """
    {
          "comment_create_time" : "2020-01-15T20:58:00+08:00",
          "like_count" : 1,
          "sub_comment_count" : 1,
          "create_time" : "2020-02-27T05:16:11",
          "avatar" : "https://img.xiaohongshu.com/avatar/5e10a13b9b48a000011ea178.jpg@80w_80h_90q_1e_1c_1x.jpg",
          "comment_id" : "5e1f0c7c000000000100a760",
          "content" : "健康[偷笑R]",
          "post_id" : "5e1e6116000000000100b1fa",
          "user_id" : "5502cf124fac6302a380c747",
          "nickname" : "邵筅笙"
        }
    """
    index = 'comment_xiaohongshu'


class ESFansRed(ESZoo):
    index = 'fans_xiaohongshu'


# 微博
class ESPostWb(ESZoo):
    index = 'post_weibo'


class ESUserWb(ESZoo):
    index = 'user_weibo'


class ESFansWb(ESZoo):
    index = 'fans_weibo'


class ESCommentWb(ESZoo):
    """
    {
          "comment_create_time" : "2019-01-28T21:48:23+08:00",
          "post_create_time" : "2019-01-28T19:24:13+08:00",
          "create_time" : "2020-01-07T17:16:09",
          "disable_reply" : 0,
          "floor_number" : 51,
          "comment_id" : "4333554227781527",
          "post_id" : "4333517942394999",
          "root_id" : "4333554227781527",
          "text" : "这么棒的创意是谁的呀，不错[可爱]",
          "user" : {
            "province" : "江西",
            "user_id" : "6678358837",
            "nickname" : "戎访蕊",
            "region" : "江西"
          },
          "readtimetype" : "comment"
        }
      }
    """
    index = 'comment_weibo'


# Bilibili
class ESUserBili(ESZoo):
    index = 'user_bsite'


class ESPostBili(ESZoo):
    index = 'post_bsite'


class ESFansBili(ESZoo):
    index = 'fans_bsite'


class ESCommentBili(ESZoo):
    """
    {
          "comment_create_time" : 1547001786,
          "liked_count" : 1,
          "post_create_time" : "2019-01-09T10:26:32+08:00",
          "create_time" : "2020-03-01T14:36:35",
          "avatar" : "http://i0.hdslb.com/bfs/face/08b32a678b36ee0baa3c28a5d2210b010286749f.jpg",
          "comment_id" : "1291485603",
          "type" : 1,
          "content" : "吕儿的口胡让阿芬送鸭肉西施了。",
          "post_id" : "40284392",
          "user_id" : "234315557",
          "nickname" : "只是风1"
        }
    """
    index = 'comment_bsite-v1_0-2018*,comment_bsite-v1_0-2019*,comment_bsite-v1_0-2020*'


# 舆情
class ESYqPost(ESYq):
    index = 'post'

#微信
class ESPostWx(ESZoo):
    index = 'post_weixin'

class ESUserWx(ESZoo):
    index = 'user_weixin'


#头条
class ESPostTt(ESZoo):
    index = 'post_toutiao'

#知乎
class ESPostZh(ESZoo):
    index = 'post_zhihu'

class ESUserZh(ESZoo):
    index = 'user_zhihu'


#天猫评论
class ESCommentTm(ESZoo):
    index = 'comment_tmall'

#天猫
class ESShopTm(ESZoo):
    index = 'shop_tmall'


# 京东评论
class ESCommentJd(ESZoo):
    index = 'comment_jd'


# 京东
class ESShopJd(ESZoo):
    index = 'shop_jd'

# 小红书logs_sku
class ESLogsSkuRed(ESZoo):
    index = 'logs_sku_xiaohongshu_2020'

# 小红书sku
class ESSkuRed(ESZoo):
    index = 'sku_xiaohongshu'